use ring::aead::{Aad, LessSafeKey, Nonce, UnboundKey, AES_256_GCM, NONCE_LEN};
use ring::rand::{SecureRandom, SystemRandom};
use ring::pbkdf2;
use std::num::NonZeroU32;
use anyhow::{Result, anyhow};
use serde::{Deserialize, Serialize};
use std::path::{Path, PathBuf};
use tokio::fs;
use log::{info, error, debug, warn};
use std::collections::HashMap;
use uuid::Uuid;
use base64::{Engine as _, engine::general_purpose};
use zeroize::Zeroize;

/// 加密的数据结构
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EncryptedData {
    /// 加密算法标识
    pub algorithm: String,
    /// 盐值（用于密钥派生）
    pub salt: Vec<u8>,
    /// 随机数
    pub nonce: Vec<u8>,
    /// 加密后的数据
    pub ciphertext: Vec<u8>,
    /// 认证标签
    pub tag: Vec<u8>,
}

/// 安全的字符串，在内存中自动清零
#[derive(Clone)]
pub struct SecureString {
    data: String,
}

impl SecureString {
    pub fn new(data: String) -> Self {
        Self { data }
    }

    pub fn as_str(&self) -> &str {
        &self.data
    }

    pub fn as_bytes(&self) -> &[u8] {
        self.data.as_bytes()
    }

    pub fn len(&self) -> usize {
        self.data.len()
    }

    pub fn is_empty(&self) -> bool {
        self.data.is_empty()
    }
}

impl std::fmt::Debug for SecureString {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "SecureString([REDACTED])")
    }
}

impl From<String> for SecureString {
    fn from(data: String) -> Self {
        Self::new(data)
    }
}

impl From<&str> for SecureString {
    fn from(data: &str) -> Self {
        Self::new(data.to_string())
    }
}

/// 密钥管理器
pub struct KeyManager {
    master_key: Option<SecureString>,
    key_store: HashMap<String, EncryptedData>,
    key_store_path: PathBuf,
    rng: SystemRandom,
}

impl KeyManager {
    /// 创建新的密钥管理器
    pub fn new(key_store_path: PathBuf) -> Self {
        Self {
            master_key: None,
            key_store: HashMap::new(),
            key_store_path,
            rng: SystemRandom::new(),
        }
    }

    /// 设置主密钥
    pub fn set_master_key(&mut self, password: SecureString) -> Result<()> {
        if password.is_empty() {
            return Err(anyhow!("主密钥不能为空"));
        }
        
        self.master_key = Some(password);
        info!("主密钥已设置");
        Ok(())
    }

    /// 检查是否已设置主密钥
    pub fn has_master_key(&self) -> bool {
        self.master_key.is_some()
    }

    /// 派生加密密钥
    fn derive_key(&self, salt: &[u8]) -> Result<LessSafeKey> {
        let master_key = self.master_key.as_ref()
            .ok_or_else(|| anyhow!("未设置主密钥"))?;
        
        let mut key_bytes = [0u8; 32]; // AES-256需要32字节密钥
        
        pbkdf2::derive(
            pbkdf2::PBKDF2_HMAC_SHA256,
            NonZeroU32::new(100_000).unwrap(), // 迭代次数
            salt,
            master_key.as_bytes(),
            &mut key_bytes,
        );
        
        let unbound_key = UnboundKey::new(&AES_256_GCM, &key_bytes)
            .map_err(|e| anyhow!("创建密钥失败: {:?}", e))?;
        
        Ok(LessSafeKey::new(unbound_key))
    }

    /// 加密数据
    pub fn encrypt(&self, plaintext: &[u8]) -> Result<EncryptedData> {
        if !self.has_master_key() {
            return Err(anyhow!("未设置主密钥"));
        }
        
        // 生成随机盐值
        let mut salt = [0u8; 32];
        self.rng.fill(&mut salt)
            .map_err(|e| anyhow!("生成盐值失败: {:?}", e))?;
        
        // 生成随机nonce
        let mut nonce_bytes = [0u8; NONCE_LEN];
        self.rng.fill(&mut nonce_bytes)
            .map_err(|e| anyhow!("生成nonce失败: {:?}", e))?;
        
        let nonce = Nonce::assume_unique_for_key(nonce_bytes);
        
        // 派生密钥
        let key = self.derive_key(&salt)?;
        
        // 加密数据
        let mut ciphertext = plaintext.to_vec();
        let tag = key.seal_in_place_separate_tag(
            nonce,
            Aad::empty(),
            &mut ciphertext,
        ).map_err(|e| anyhow!("加密失败: {:?}", e))?;
        
        Ok(EncryptedData {
            algorithm: "AES-256-GCM".to_string(),
            salt: salt.to_vec(),
            nonce: nonce_bytes.to_vec(),
            ciphertext,
            tag: tag.as_ref().to_vec(),
        })
    }

    /// 解密数据
    pub fn decrypt(&self, encrypted_data: &EncryptedData) -> Result<Vec<u8>> {
        if !self.has_master_key() {
            return Err(anyhow!("未设置主密钥"));
        }
        
        if encrypted_data.algorithm != "AES-256-GCM" {
            return Err(anyhow!("不支持的加密算法: {}", encrypted_data.algorithm));
        }
        
        // 派生密钥
        let key = self.derive_key(&encrypted_data.salt)?;
        
        // 重构nonce
        if encrypted_data.nonce.len() != NONCE_LEN {
            return Err(anyhow!("无效的nonce长度"));
        }
        
        let mut nonce_bytes = [0u8; NONCE_LEN];
        nonce_bytes.copy_from_slice(&encrypted_data.nonce);
        let nonce = Nonce::assume_unique_for_key(nonce_bytes);
        
        // 准备解密数据
        let mut ciphertext_and_tag = encrypted_data.ciphertext.clone();
        ciphertext_and_tag.extend_from_slice(&encrypted_data.tag);
        
        // 解密
        let plaintext = key.open_in_place(
            nonce,
            Aad::empty(),
            &mut ciphertext_and_tag,
        ).map_err(|e| anyhow!("解密失败: {:?}", e))?;
        
        Ok(plaintext.to_vec())
    }

    /// 存储加密的密钥
    pub async fn store_key(&mut self, key_id: &str, key_data: &[u8]) -> Result<()> {
        let encrypted_data = self.encrypt(key_data)?;
        self.key_store.insert(key_id.to_string(), encrypted_data);
        
        // 保存到文件
        self.save_key_store().await?;
        
        info!("存储密钥: {}", key_id);
        Ok(())
    }

    /// 检索解密的密钥
    pub fn retrieve_key(&self, key_id: &str) -> Result<Vec<u8>> {
        let encrypted_data = self.key_store.get(key_id)
            .ok_or_else(|| anyhow!("密钥不存在: {}", key_id))?;
        
        self.decrypt(encrypted_data)
    }

    /// 删除密钥
    pub async fn delete_key(&mut self, key_id: &str) -> Result<()> {
        if self.key_store.remove(key_id).is_some() {
            self.save_key_store().await?;
            info!("删除密钥: {}", key_id);
        }
        Ok(())
    }

    /// 列出所有密钥ID
    pub fn list_keys(&self) -> Vec<String> {
        self.key_store.keys().cloned().collect()
    }

    /// 保存密钥存储到文件
    async fn save_key_store(&self) -> Result<()> {
        let content = serde_json::to_string_pretty(&self.key_store)
            .map_err(|e| anyhow!("序列化密钥存储失败: {}", e))?;
        
        // 确保目录存在
        if let Some(parent) = self.key_store_path.parent() {
            fs::create_dir_all(parent).await
                .map_err(|e| anyhow!("创建密钥存储目录失败: {}", e))?;
        }
        
        fs::write(&self.key_store_path, content).await
            .map_err(|e| anyhow!("保存密钥存储失败: {}", e))?;
        
        debug!("密钥存储已保存: {:?}", self.key_store_path);
        Ok(())
    }

    /// 从文件加载密钥存储
    pub async fn load_key_store(&mut self) -> Result<()> {
        if !self.key_store_path.exists() {
            debug!("密钥存储文件不存在，创建新的存储");
            return Ok(());
        }
        
        let content = fs::read_to_string(&self.key_store_path).await
            .map_err(|e| anyhow!("读取密钥存储失败: {}", e))?;
        
        self.key_store = serde_json::from_str(&content)
            .map_err(|e| anyhow!("解析密钥存储失败: {}", e))?;
        
        info!("加载了 {} 个密钥", self.key_store.len());
        Ok(())
    }

    /// 更改主密钥
    pub async fn change_master_key(&mut self, old_password: SecureString, new_password: SecureString) -> Result<()> {
        // 验证旧密钥
        let old_master_key = self.master_key.clone();
        self.master_key = Some(old_password);
        
        // 尝试解密一个已存储的密钥来验证旧密钥
        if !self.key_store.is_empty() {
            let test_key_id = self.key_store.keys().next().unwrap();
            if let Err(e) = self.retrieve_key(test_key_id) {
                self.master_key = old_master_key;
                return Err(anyhow!("旧密钥验证失败: {}", e));
            }
        }
        
        // 解密所有密钥
        let mut decrypted_keys = HashMap::new();
        for (key_id, encrypted_data) in &self.key_store {
            let decrypted = self.decrypt(encrypted_data)?;
            decrypted_keys.insert(key_id.clone(), decrypted);
        }
        
        // 设置新密钥
        self.master_key = Some(new_password);
        
        // 用新密钥重新加密所有密钥
        self.key_store.clear();
        for (key_id, key_data) in decrypted_keys {
            let encrypted_data = self.encrypt(&key_data)?;
            self.key_store.insert(key_id, encrypted_data);
        }
        
        // 保存更新后的密钥存储
        self.save_key_store().await?;
        
        info!("主密钥已更改");
        Ok(())
    }
}

/// 密码强度检查器
pub struct PasswordStrengthChecker;

impl PasswordStrengthChecker {
    /// 检查密码强度
    pub fn check_strength(password: &str) -> PasswordStrength {
        let mut score: i32 = 0;
        let mut feedback = Vec::new();
        
        // 长度检查
        if password.len() >= 12 {
            score += 2;
        } else if password.len() >= 8 {
            score += 1;
        } else {
            feedback.push("密码长度至少应为8个字符".to_string());
        }
        
        // 字符类型检查
        let has_lowercase = password.chars().any(|c| c.is_ascii_lowercase());
        let has_uppercase = password.chars().any(|c| c.is_ascii_uppercase());
        let has_digit = password.chars().any(|c| c.is_ascii_digit());
        let has_special = password.chars().any(|c| !c.is_alphanumeric());
        
        if has_lowercase { score += 1; }
        if has_uppercase { score += 1; }
        if has_digit { score += 1; }
        if has_special { score += 1; }
        
        if !has_lowercase { feedback.push("应包含小写字母".to_string()); }
        if !has_uppercase { feedback.push("应包含大写字母".to_string()); }
        if !has_digit { feedback.push("应包含数字".to_string()); }
        if !has_special { feedback.push("应包含特殊字符".to_string()); }
        
        // 重复字符检查
        let mut char_counts = HashMap::new();
        for c in password.chars() {
            *char_counts.entry(c).or_insert(0) += 1;
        }
        
        let max_repeat = char_counts.values().max().unwrap_or(&0);
        if *max_repeat > 3 {
            feedback.push("避免重复字符过多".to_string());
            score = score.saturating_sub(1);
        }
        
        // 常见密码检查
        let common_passwords = [
            "password", "123456", "123456789", "qwerty", "abc123",
            "password123", "admin", "letmein", "welcome", "monkey"
        ];
        
        if common_passwords.iter().any(|&p| password.to_lowercase().contains(p)) {
            feedback.push("避免使用常见密码".to_string());
            score = score.saturating_sub(2);
        }
        
        let level = match score {
            0..=2 => PasswordStrengthLevel::Weak,
            3..=4 => PasswordStrengthLevel::Fair,
            5..=6 => PasswordStrengthLevel::Good,
            _ => PasswordStrengthLevel::Strong,
        };
        
        PasswordStrength {
            level,
            score: score as u32,
            feedback,
        }
    }
}

/// 密码强度等级
#[derive(Debug, Clone, PartialEq)]
pub enum PasswordStrengthLevel {
    Weak,
    Fair,
    Good,
    Strong,
}

/// 密码强度结果
#[derive(Debug, Clone)]
pub struct PasswordStrength {
    pub level: PasswordStrengthLevel,
    pub score: u32,
    pub feedback: Vec<String>,
}

/// SSH密钥验证器
pub struct SshKeyValidator;

impl SshKeyValidator {
    /// 验证SSH私钥文件
    pub async fn validate_private_key(key_path: &Path) -> Result<SshKeyInfo> {
        if !key_path.exists() {
            return Err(anyhow!("密钥文件不存在: {:?}", key_path));
        }
        
        let content = fs::read_to_string(key_path).await
            .map_err(|e| anyhow!("读取密钥文件失败: {}", e))?;
        
        // 检查文件权限（Unix系统）
        #[cfg(unix)]
        {
            use std::os::unix::fs::PermissionsExt;
            let metadata = fs::metadata(key_path).await
                .map_err(|e| anyhow!("获取文件元数据失败: {}", e))?;
            let permissions = metadata.permissions().mode() & 0o777;
            
            if permissions != 0o600 {
                warn!("SSH密钥文件权限不安全: {:o}，建议设置为600", permissions);
            }
        }
        
        // 检测密钥类型
        let key_type = if content.contains("BEGIN RSA PRIVATE KEY") {
            SshKeyType::Rsa
        } else if content.contains("BEGIN DSA PRIVATE KEY") {
            SshKeyType::Dsa
        } else if content.contains("BEGIN EC PRIVATE KEY") {
            SshKeyType::Ecdsa
        } else if content.contains("BEGIN OPENSSH PRIVATE KEY") {
            SshKeyType::Ed25519
        } else if content.contains("BEGIN PRIVATE KEY") {
            SshKeyType::Pkcs8
        } else {
            return Err(anyhow!("无法识别的SSH密钥格式"));
        };
        
        // 检查是否加密
        let is_encrypted = content.contains("Proc-Type: 4,ENCRYPTED") ||
                          content.contains("DEK-Info:");
        
        Ok(SshKeyInfo {
            key_type,
            is_encrypted,
            file_path: key_path.to_path_buf(),
        })
    }

    /// 验证SSH公钥文件
    pub async fn validate_public_key(key_path: &Path) -> Result<SshPublicKeyInfo> {
        if !key_path.exists() {
            return Err(anyhow!("公钥文件不存在: {:?}", key_path));
        }
        
        let content = fs::read_to_string(key_path).await
            .map_err(|e| anyhow!("读取公钥文件失败: {}", e))?;
        
        let lines: Vec<&str> = content.lines().collect();
        if lines.is_empty() {
            return Err(anyhow!("公钥文件为空"));
        }
        
        let first_line = lines[0];
        let parts: Vec<&str> = first_line.split_whitespace().collect();
        
        if parts.len() < 2 {
            return Err(anyhow!("无效的公钥格式"));
        }
        
        let key_type = match parts[0] {
            "ssh-rsa" => SshKeyType::Rsa,
            "ssh-dss" => SshKeyType::Dsa,
            "ecdsa-sha2-nistp256" | "ecdsa-sha2-nistp384" | "ecdsa-sha2-nistp521" => SshKeyType::Ecdsa,
            "ssh-ed25519" => SshKeyType::Ed25519,
            _ => return Err(anyhow!("不支持的公钥类型: {}", parts[0])),
        };
        
        let comment = if parts.len() > 2 {
            Some(parts[2..].join(" "))
        } else {
            None
        };
        
        Ok(SshPublicKeyInfo {
            key_type,
            comment,
            file_path: key_path.to_path_buf(),
        })
    }
}

/// SSH密钥类型
#[derive(Debug, Clone, PartialEq)]
pub enum SshKeyType {
    Rsa,
    Dsa,
    Ecdsa,
    Ed25519,
    Pkcs8,
}

/// SSH私钥信息
#[derive(Debug, Clone)]
pub struct SshKeyInfo {
    pub key_type: SshKeyType,
    pub is_encrypted: bool,
    pub file_path: PathBuf,
}

/// SSH公钥信息
#[derive(Debug, Clone)]
pub struct SshPublicKeyInfo {
    pub key_type: SshKeyType,
    pub comment: Option<String>,
    pub file_path: PathBuf,
}

/// 安全工具集合
pub struct SecurityUtils;

impl SecurityUtils {
    /// 生成安全的随机密码
    pub fn generate_password(length: usize, include_symbols: bool) -> Result<String> {
        if length < 4 {
            return Err(anyhow!("密码长度至少为4个字符"));
        }
        
        let rng = SystemRandom::new();
        let mut password = String::new();
        
        let lowercase = "abcdefghijklmnopqrstuvwxyz";
        let uppercase = "ABCDEFGHIJKLMNOPQRSTUVWXYZ";
        let digits = "0123456789";
        let symbols = if include_symbols { "!@#$%^&*()_+-=[]{}|;:,.<>?" } else { "" };
        
        let charset = format!("{}{}{}{}", lowercase, uppercase, digits, symbols);
        let charset_bytes = charset.as_bytes();
        
        for _ in 0..length {
            let mut random_byte = [0u8; 1];
            rng.fill(&mut random_byte)
                .map_err(|e| anyhow!("生成随机数失败: {:?}", e))?;
            
            let index = (random_byte[0] as usize) % charset_bytes.len();
            password.push(charset_bytes[index] as char);
        }
        
        Ok(password)
    }

    /// 计算数据的哈希值
    pub fn hash_data(data: &[u8]) -> String {
        use ring::digest;
        let digest = digest::digest(&digest::SHA256, data);
        general_purpose::STANDARD.encode(digest.as_ref())
    }

    /// 安全比较两个字符串（防止时序攻击）
    pub fn secure_compare(a: &str, b: &str) -> bool {
        if a.len() != b.len() {
            return false;
        }
        // 使用简单的字节比较，在实际应用中可以考虑使用专门的常量时间比较库
        a.as_bytes() == b.as_bytes()
    }

    /// 清理敏感数据
    pub fn clear_sensitive_data(data: &mut [u8]) {
        data.zeroize();
    }
}

/// 便捷函数：创建密钥管理器
pub async fn create_key_manager(config_dir: &Path) -> Result<KeyManager> {
    let key_store_path = config_dir.join("keys.json");
    let mut manager = KeyManager::new(key_store_path);
    manager.load_key_store().await?;
    Ok(manager)
}

/// 便捷函数：加密字符串
pub fn encrypt_string(key_manager: &KeyManager, plaintext: &str) -> Result<String> {
    let encrypted_data = key_manager.encrypt(plaintext.as_bytes())?;
    let json = serde_json::to_string(&encrypted_data)
        .map_err(|e| anyhow!("序列化加密数据失败: {}", e))?;
    Ok(general_purpose::STANDARD.encode(json))
}

/// 便捷函数：解密字符串
pub fn decrypt_string(key_manager: &KeyManager, ciphertext: &str) -> Result<String> {
    let json_bytes = general_purpose::STANDARD.decode(ciphertext)
        .map_err(|e| anyhow!("Base64解码失败: {}", e))?;
    let json = String::from_utf8(json_bytes)
        .map_err(|e| anyhow!("UTF-8解码失败: {}", e))?;
    let encrypted_data: EncryptedData = serde_json::from_str(&json)
        .map_err(|e| anyhow!("反序列化加密数据失败: {}", e))?;
    let plaintext_bytes = key_manager.decrypt(&encrypted_data)?;
    String::from_utf8(plaintext_bytes)
        .map_err(|e| anyhow!("UTF-8解码失败: {}", e))
}